/*
 * LdapExtLoginModuleForest.java
 * (c) 2011-2012 Forgotten Lian
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
 */
package org.jboss.security.auth.spi;

import java.util.Map;
import java.util.Properties;
import java.util.StringTokenizer;
import javax.naming.Context;
import javax.naming.NamingException;
import javax.naming.NamingEnumeration;
import javax.naming.ldap.InitialLdapContext;
import javax.naming.directory.Attribute;
import javax.naming.directory.Attributes;
import javax.naming.directory.SearchControls;
import javax.naming.directory.SearchResult;
import javax.security.auth.Subject;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.login.LoginException;
import org.jboss.security.util.StringPropertyReplacer;

/**
 * The LdapExtLoginModuleForest extend LdapExtLoginModule 
 * for reading user's distinguished name from the subLDAP where his record kept.
 *
 * <p>An user logins to LDAP with name is xxx@yyy.zzz
 * It is the UPN (userPrincipalName)</p>
 *
 * <ul>
 * <li><strong>rootDN</strong> - Base domain component.</li>
 * <li><strong>subDNProperties</strong> - Name of properties file. His format is domain=provider_url</li>
 * </ul>
 *
 * <p>If user's domain (yyy.zzz) is rootDN, then he has standard (LdapExtLoginModule) login procedure.</p>
 * <p>If user is from subdomain (vv.yyy.zzz), then LoginModule:
 * <ul>
 * <li>does bind to the user's LDAP server</li>
 * <li>gets user's Distinguished Name and return it to LdapLoginModule</li>
 * </ul>
 * </p>
 *
 * @author Andrey Kotov
 * @version 1.2
 */
public class LdapExtLoginModuleForest extends LdapExtLoginModule {

    private static final String DISTINGUISHED_NAME_ATTRIBUTE_OPT = "distinguishedNameAttribute";
    private static final String ROOT_DN_OPT = "rootDN"; // Base domain components e.g. mycompany.com, where lookup roles
    private static final String SUB_DN_PROPERTIES_OPT = "subDNProperties"; // Name of properties file contains ldaps addresses
    protected String rootDN;
    protected String subDNProperties = "subLdaps.properties";
    protected Properties ldapForest;

    public LdapExtLoginModuleForest() {
    }

    @Override
    @SuppressWarnings("unchecked")
    public void initialize(Subject subject, CallbackHandler callbackHandler, Map sharedState, Map options) {
        super.initialize(subject, callbackHandler, sharedState, options);

        trace = log.isTraceEnabled();

        String opt = (String) options.get(ROOT_DN_OPT);
        if (opt != null) {
            rootDN = opt;
        }
        if (trace) {
            log.trace("Base domain: " + (rootDN == null ? "null" : rootDN));
        }

        opt = (String) options.get(SUB_DN_PROPERTIES_OPT);
        if (opt != null) {
            subDNProperties = StringPropertyReplacer.replaceProperties(opt);
        }
        if (trace) {
            log.trace("Name properties file with ldap: " + subDNProperties);
        }

        opt = (String) options.get(DISTINGUISHED_NAME_ATTRIBUTE_OPT);
        if (opt != null) {
            distinguishedNameAttribute = opt;
        }

        try {
            ldapForest = Util.loadProperties(subDNProperties, log);
        } catch (Exception e) {
            log.error("Failed to load ldaps properties files", e);
        }
    }

    @Override
    public boolean login() throws LoginException {
        if (ldapForest == null) {
            throw new LoginException("Missing " + subDNProperties + " file.");
        }
        return super.login();
    }

    @Override
    protected String bindDNAuthentication(InitialLdapContext ctx, String user, Object credential, String baseDN, String filter) throws NamingException {
        String userDC = null;
        int pos = user.indexOf('@');
        if (pos > 0) {
            userDC = user.substring(pos + 1);
        }
        if (trace) {
            log.trace("user domain = '" + (userDC == null ? "null" : userDC) + "'");
        }

        // Default bind
        if (userDC == null || userDC.equals(rootDN)) {
            if (trace) {
                log.trace("Get userDN wich standard way");
            }
            return super.bindDNAuthentication(ctx, user, credential, baseDN, filter);
        }

        // Get user's ldap
        String userLDAP = ldapForest.getProperty(userDC);
        if (userLDAP == null) {
            throw new NamingException("For user '" + user + "' not found ldap for user domain '" + userDC + "' in " + subDNProperties);
        }

        // Construct userCtxDN
        StringBuilder sbUserCtxDN = new StringBuilder();
        StringTokenizer st = new StringTokenizer(userDC, ".");
        while (st.hasMoreTokens()) {
            if (sbUserCtxDN.length() > 0) {
                sbUserCtxDN.append(",");
            }
            sbUserCtxDN.append("dc=");
            sbUserCtxDN.append(st.nextToken());
        }
        String userCtxDN = sbUserCtxDN.toString(); // user's baseDN
        if (trace) {
            log.trace("user context domain = '" + (userCtxDN == null ? "null" : userCtxDN) + "'");
        }

        // Create the user context for searching distinguished name
        Properties env = new Properties();
        String opt = (String) options.get(Context.INITIAL_CONTEXT_FACTORY);
        if (opt == null) {
            opt = "com.sun.jndi.ldap.LdapCtxFactory";
        }
        env.setProperty(Context.INITIAL_CONTEXT_FACTORY, opt);

        env.setProperty(Context.PROVIDER_URL, userLDAP);

        opt = (String) options.get(Context.SECURITY_AUTHENTICATION);
        if (opt == null) {
            opt = "simple";
        }
        env.setProperty(Context.SECURITY_AUTHENTICATION, opt);

        env.setProperty(Context.SECURITY_PRINCIPAL, user);
        env.setProperty(Context.SECURITY_CREDENTIALS, (String) credential);

        if (trace) {
            Properties tmp = new Properties();
            tmp.putAll(env);
            tmp.setProperty(Context.SECURITY_CREDENTIALS, "***");
            log.trace("Logging into USER LDAP server, env=" + tmp.toString());
        }

        InitialLdapContext userCtx = new InitialLdapContext(env, null);

        // Search the distinguished name in the user's ldap
        SearchControls constraints = new SearchControls();
        constraints.setSearchScope(SearchControls.SUBTREE_SCOPE);
        constraints.setTimeLimit(searchTimeLimit);
        String attrList[] = {distinguishedNameAttribute};
        constraints.setReturningAttributes(attrList);

        NamingEnumeration results = null;

        Object[] filterArgs = {user};
        results = userCtx.search(userCtxDN, filter, filterArgs, constraints);
        if (results.hasMore() == false) {
            results.close();
            throw new NamingException("Search of userCtxDN(" + userCtxDN + ") found no matches");
        }

        SearchResult sr = (SearchResult) results.next();
        String name = sr.getName();
        String userDN = null;
        Attributes attrs = sr.getAttributes();
        if (attrs != null) {
            Attribute dn = attrs.get(distinguishedNameAttribute);
            if (dn != null) {
                userDN = (String) dn.get();
                if (trace) {
                    log.trace("userDN (distinguished name) = " + userDN);
                }
            }
        }
        if (userDN == null) {
            if (sr.isRelative() == true) {
                userDN = name + "," + userCtxDN;
                if (trace) {
                    log.trace("userDN (relative distinguished name) = " + userDN);
                }
            } else {
                throw new NamingException("Can't follow referal for authentication: " + name);
            }
        }

        results.close();
        results = null;
        userCtx.close();

        return userDN;
    }
}
